# 首先要安装LlamIndex  pip install llama-index
from llama_index.core import SimpleDirectoryReader, Settings, VectorStoreIndex
from embeddings import embed_model_local_bge_small
from llm import deepseek_llm,openAi_llm
import os

root_dir = os.path.dirname(os.path.abspath(__file__))

# 加载指定目录下的文件
documents = SimpleDirectoryReader(root_dir +'/data').load_data()

# 设置embedding模型
Settings.embed_model = embed_model_local_bge_small()

# 设置LLM模型
Settings.llm = openAi_llm()

#创建向量索引
index = VectorStoreIndex.from_documents(
    documents,
)

query_engine = index.as_query_engine()
response = query_engine.query("员工年假？")
print(response)